import numpy as np
from loadGripper import get_two_gripper_pcd
from loadFinger import get_4finger_from_gripper36
MIMIC_GRIPPER = False
def gripper(unified_data: dict, gripper_key: str) -> dict:
    '''
    # ONLY USEFUL WHEN CHOOSING EE ACTION TYPE. USING 'state' AS EE OBS
    '''
    # Check for repointcloud and gc
    repointcloud = unified_data.get('repointcloud', None)
    gc = unified_data.get('gc', None)
    state = unified_data.get('state', None)
    assert repointcloud is not None, "repointcloud is None in unified_data"
    assert repointcloud is not False, "repointcloud is False in unified_data"
    assert gc is not None, "gc is None in unified_data"
    assert gc is not False, "gc is False in unified_data"
    assert state is not None, "state is None in unified_data"
    assert len(state) == 16, f"state length is not 16 in unified_data {unified_data['state'].shape}"
    num_points = int(gc) if np.isscalar(gc) else int(gc[()])
    assert num_points > 0, "num_points is not positive"
    pose1 = state[0:7]
    pose2 = state[8:15]
    # Extract finger_dist from state (assuming it's the last element of each arm's state)
    finger_dist1 = state[7]  # Left gripper finger distance
    finger_dist2 = state[15]  # Right gripper finger distance
    # Calculate point distribution: 10% for 4 fingers, 90% for 2 grippers
    n_gripper_total = int(num_points * 0.9)  # 90% for grippers
    n_finger_total = num_points - n_gripper_total  # 10% for fingers
    
    print("state",list(np.round(np.array(state),2)))
    print("pose1",list(np.round(pose1,2)))
    print("pose2",list(np.round(pose2,2)))
    print("finger_dist1", finger_dist1, "finger_dist2", finger_dist2)
    print(f"Total points: {num_points}, Gripper points: {n_gripper_total}, Finger points: {n_finger_total}")
    
    # Generate gripper point clouds (90% of total points)
    gripper_pcds = get_two_gripper_pcd(gripper_key, n_gripper_total, pose1, pose2, finger_dist=finger_dist1)
    assert (len(gripper_pcds) > 0), f"Failed to generate real gripper point clouds {gripper_key} due to pcd1:{len(gripper_pcds)} points, falling back to rectangle generation"
    
    # Generate finger point clouds (10% of total points)
    # Get gripper36 data from key_poses [12:48] in unified_data
    gripper36 = unified_data['key_poses'][12:48]  # 36D gripper pose data
    
    # Generate finger point clouds
    finger_pcds = get_4finger_from_gripper36(gripper_key, n_finger_total, gripper36)
    
    # Combine gripper and finger point clouds
    combined_pcds = np.concatenate([gripper_pcds, finger_pcds], axis=0)
    print(f"Successfully generated {len(gripper_pcds)} gripper points and {len(finger_pcds)} finger points")
    
    unified_data['repointcloud'] = np.concatenate([repointcloud, combined_pcds], axis=0)
    return unified_data 